package com.eian.boot.crypto.store;

import com.eian.boot.crypto.constant.CryptoConstants;
import com.eian.boot.crypto.model.CryptoContext;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 基于内存的密钥存储实现
 * <p>
 * 使用 ConcurrentHashMap 存储密钥信息，支持并发访问
 * <p>
 * 默认配置：
 * <ul>
 *     <li>最大存储数量：1000</li>
 *     <li>密钥过期时间：3600秒（1小时）</li>
 * </ul>
 * <p>
 * 适用场景：单机环境或不需要密钥持久化的场景
 *
 * @author eian
 */
@Slf4j
public class InMemoryKeyStore implements KeyStore {

    /**
     * 密钥存储容器
     */
    private final Map<String, KeyEntry> storage = new ConcurrentHashMap<>();

    /**
     * 最大存储数量
     */
    private final int maxSize;

    /**
     * 密钥过期时间（秒）
     */
    private final long expireSeconds;

    /**
     * 定时清理任务
     */
    private final ScheduledExecutorService cleanupExecutor;

    /**
     * 使用默认配置构造
     */
    public InMemoryKeyStore() {
        this(CryptoConstants.KeyStore.DEFAULT_MAX_SIZE, CryptoConstants.KeyStore.DEFAULT_EXPIRE_SECONDS);
    }

    /**
     * 自定义配置构造
     *
     * @param maxSize       最大存储数量
     * @param expireSeconds 过期时间（秒，0表示不过期）
     */
    public InMemoryKeyStore(int maxSize, long expireSeconds) {
        this.maxSize = maxSize;
        this.expireSeconds = expireSeconds;

        // 如果配置了过期时间，启动定时清理任务
        if (expireSeconds > 0) {
            this.cleanupExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
                Thread thread = new Thread(r, "KeyStore-Cleanup");
                thread.setDaemon(true);
                return thread;
            });
            // 每5分钟执行一次清理
            cleanupExecutor.scheduleAtFixedRate(this::cleanupExpired, 5, 5, TimeUnit.MINUTES);
            log.info("InMemoryKeyStore initialized with maxSize={}, expireSeconds={}", maxSize, expireSeconds);
        } else {
            this.cleanupExecutor = null;
            log.info("InMemoryKeyStore initialized with maxSize={}, no expiration", maxSize);
        }
    }

    @Override
    public void store(String keyId, CryptoContext cryptoContext) {
        if (keyId == null || cryptoContext == null) {
            log.warn("Cannot store null keyId or cryptoContext");
            return;
        }

        // 检查容量限制
        if (storage.size() >= maxSize && !storage.containsKey(keyId)) {
            log.warn("KeyStore reached max size {}, cannot store new key: {}", maxSize, keyId);
            return;
        }

        long expireTime = expireSeconds > 0
                ? System.currentTimeMillis() + expireSeconds * 1000
                : 0;

        storage.put(keyId, new KeyEntry(cryptoContext, expireTime));
        log.debug("Stored key: {}, current size: {}", keyId, storage.size());
    }

    @Override
    public CryptoContext retrieve(String keyId) {
        if (keyId == null) {
            return null;
        }

        KeyEntry entry = storage.get(keyId);
        if (entry == null) {
            return null;
        }

        // 检查是否过期
        if (entry.isExpired()) {
            storage.remove(keyId);
            log.debug("Key expired and removed: {}", keyId);
            return null;
        }

        return entry.getContext();
    }

    @Override
    public boolean remove(String keyId) {
        if (keyId == null) {
            return false;
        }

        boolean removed = storage.remove(keyId) != null;
        if (removed) {
            log.debug("Removed key: {}, current size: {}", keyId, storage.size());
        }
        return removed;
    }

    @Override
    public boolean exists(String keyId) {
        if (keyId == null) {
            return false;
        }

        KeyEntry entry = storage.get(keyId);
        if (entry == null) {
            return false;
        }

        // 检查是否过期
        if (entry.isExpired()) {
            storage.remove(keyId);
            return false;
        }

        return true;
    }

    @Override
    public void clear() {
        int size = storage.size();
        storage.clear();
        log.info("Cleared all keys, removed {} entries", size);
    }

    /**
     * 清理过期的密钥
     */
    private void cleanupExpired() {
        if (expireSeconds <= 0) {
            return;
        }

        int removed = 0;
        long now = System.currentTimeMillis();

        for (Map.Entry<String, KeyEntry> entry : storage.entrySet()) {
            if (entry.getValue().isExpiredAt(now)) {
                storage.remove(entry.getKey());
                removed++;
            }
        }

        if (removed > 0) {
            log.info("Cleaned up {} expired keys, current size: {}", removed, storage.size());
        }
    }

    /**
     * 获取当前存储的密钥数量
     */
    public int size() {
        return storage.size();
    }

    /**
     * 关闭密钥存储（关闭清理任务）
     */
    public void shutdown() {
        if (cleanupExecutor != null && !cleanupExecutor.isShutdown()) {
            cleanupExecutor.shutdown();
            log.info("KeyStore cleanup executor shutdown");
        }
    }

    /**
     * 密钥条目（包含过期时间）
     */
    private static class KeyEntry {
        private final CryptoContext context;
        private final long expireTime;

        KeyEntry(CryptoContext context, long expireTime) {
            this.context = context;
            this.expireTime = expireTime;
        }

        CryptoContext getContext() {
            return context;
        }

        boolean isExpired() {
            return expireTime > 0 && System.currentTimeMillis() > expireTime;
        }

        boolean isExpiredAt(long timestamp) {
            return expireTime > 0 && timestamp > expireTime;
        }
    }
}

